package com.mc.book.utils.baidu;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.mc.book.utils.sse.BookListener;
import com.mc.book.utils.sse.ExecuteSSEUtil;
import okhttp3.*;
import org.springframework.web.context.request.RequestContextHolder;

import javax.servlet.http.HttpServletResponse;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.net.HttpURLConnection;
import java.net.URL;


public class WenXinUtil {
    static final OkHttpClient HTTP_CLIENT = new OkHttpClient();

    static String client_id = "RkAXvxGyT2mn9soKAjLNOLo9";
    static String client_secret = "9AP6lwvpG6KLNtuIk2Me4wwhMvppAqqA";

    public static void main(String[] args) {
//        System.out.println(chat("您是擅长讲解图书的故事高手，我想学习《论语》这本书的精华内容，请您结合故事，使用通俗易懂、风趣幽默的语言，详细解读这本书的核心内容。"));
        System.out.println(chat("hello"));
    }


    public static String getToken(){
        String url = "https://aip.baidubce.com/oauth/2.0/token";
        String grant_type = "client_credentials";

        MediaType mediaType = MediaType.parse("application/x-www-form-urlencoded");
        RequestBody body = RequestBody.create(mediaType, "grant_type=client_credentials&client_id=" + client_id
                + "&client_secret=" + client_secret);
        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/oauth/2.0/token")
                .method("POST", body)
                .addHeader("Content-Type", "application/x-www-form-urlencoded")
                .build();
        Response response = null;
        try {
            response = HTTP_CLIENT.newCall(request).execute();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        try {
            return JSONObject.parseObject(response.body().string()).getString("access_token");
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public static void chatStream(String content, HttpServletResponse response){
        String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie_speed";
        String token = getToken();

        JSONObject message = new JSONObject();
        message.put("role","user");
        message.put("content",content);

        JSONArray messages = new JSONArray();
        messages.add(message);

        JSONObject param = new JSONObject();
        param.put("messages",messages);
        param.put("stream",true);

        try {
            BookListener sseListener = new BookListener(response);
            ExecuteSSEUtil.executeSSE(url+"?access_token=" + token, sseListener, param.toString());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public static String chat(String content){
        String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie_speed";
        String token = getToken();

        JSONObject message = new JSONObject();
        message.put("role","user");
        message.put("content",content);

        JSONArray messages = new JSONArray();
        messages.add(message);

        JSONObject param = new JSONObject();
        param.put("messages",messages);

        try {
            String result = getURLByPost(url+"?access_token=" + token,param.toString());
            return JSONObject.parseObject(result).getString("result");
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public static String getURLByPost(String urlStr,String params)throws Exception{
        URL url=new URL(urlStr);
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod("POST");
        conn.setRequestProperty("Content-type", "application/json");
        conn.setDoOutput(true);
        conn.setDoInput(true);
        PrintWriter printWriter = new PrintWriter(conn.getOutputStream());
        printWriter.write(params);
        printWriter.flush();
        BufferedReader in = null;
        StringBuilder sb = new StringBuilder();
        try{
            in = new BufferedReader( new InputStreamReader(conn.getInputStream(),"UTF-8") );
            String str = null;
            while((str = in.readLine()) != null) {
                sb.append( str );
            }
        } catch (Exception ex) {
            throw ex;
        } finally{
            try{
                conn.disconnect();
                if(in!=null){
                    in.close();
                }
                if(printWriter!=null){
                    printWriter.close();
                }
            }catch(IOException ex) {
                throw ex;
            }
        }
        return sb.toString();
    }

}
